#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
完整目标检测 YOLO 节点（带结果缓存队列）

功能说明：
1. 订阅 /yolo_detect 话题接收控制命令（yolostart触发完整检测任务）
2. 订阅 /camera/image_raw 话题接收相机图像
3. 接收到yolostart命令时，开始持续检测直到三个目标都识别成功：
   - 锥形桶颜色识别
   - 数字识别
   - 仪表盘状态识别
4. 使用三个结果缓存队列管理检测结果
5. 三个目标都识别成功后，发布完整结果到/yolo_result话题并停止检测
6. 发布状态信息到/yolo_status话题

话题接口：
- 订阅：/yolo_detect (std_msgs/String) - 控制命令
- 订阅：/camera/image_raw (sensor_msgs/Image) - 相机图像  
- 发布：/yolo_result (std_msgs/String) - 检测结果
- 发布：/yolo_status (std_msgs/String) - 状态信息

检测结果格式：color-number-state (如: red-three-up)
"""

import os
import sys
import gc
import time
import threading
from collections import deque
from typing import Optional, Dict, Any

import rospy
import cv2
import numpy as np
from std_msgs.msg import String
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

# 添加yolo-jetson路径到sys.path以导入管道
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, "../../../../"))
yolo_jetson_path = os.path.join(project_root, "yolo-jetson", "detect")
if os.path.exists(yolo_jetson_path):
    sys.path.insert(0, yolo_jetson_path)

try:
    from double_yolo_pipeline import DoubleYoloPipeline
    rospy.loginfo("成功导入 DoubleYoloPipeline")
except ImportError as e:
    rospy.logerr(f"无法导入 DoubleYoloPipeline: {e}")
    rospy.logerr(f"检查路径: {yolo_jetson_path}")
    sys.exit(1)


class ResultQueue:
    """结果缓存队列类"""
    
    def __init__(self, max_size=5, name=""):
        self.queue = deque(maxlen=max_size)
        self.max_size = max_size
        self.name = name
        
    def add(self, result):
        """添加结果到队列"""
        if result is not None and result != 'None':
            self.queue.append(result)
            return True
        return False
        
    def is_full(self):
        """检查队列是否已满"""
        return len(self.queue) >= self.max_size
        
    def get_most_common(self):
        """获取队列中最常见的结果"""
        if not self.queue:
            return None
        # 简单返回最后一个结果，也可以实现投票机制
        return self.queue[-1]
    
    def clear(self):
        """清空队列"""
        self.queue.clear()
        
    def get_status(self):
        """获取队列状态"""
        return f"{self.name}: {len(self.queue)}/{self.max_size} {list(self.queue)}"


class YoloCompleteDetectionNode:
    """完整目标检测YOLO节点（带结果缓存队列）"""
    
    def __init__(self):
        rospy.init_node('yolo_complete_detection_node', anonymous=True)
        
        # 参数配置
        self.model_path = rospy.get_param('~model_path', '')
        self.device = rospy.get_param('~device', None)  # None表示自动选择
        self.debug = rospy.get_param('~debug', True)
        self.save_dir = rospy.get_param('~save_dir', '/tmp/yolo_complete_detection_debug')
        
        # 检测参数
        self.top_cut_ratio = rospy.get_param('~top_cut_ratio', 0.25)
        self.width_center_ratio = rospy.get_param('~width_center_ratio', 0.5)
        self.first_stage_conf = rospy.get_param('~first_stage_conf', 0.6)
        self.second_stage_conf = rospy.get_param('~second_stage_conf', 0.5)
        self.final_conf = rospy.get_param('~final_conf', 0.6)
        self.process_interval = rospy.get_param('~process_interval', 0.1)  # 帧处理间隔
        self.max_detection_time = rospy.get_param('~max_detection_time', 30.0)  # 最大检测时间（秒）
        
        # 结果队列配置
        self.queue_size = rospy.get_param('~queue_size', 5)  # 队列大小
        
        # 状态管理
        self.pipeline_initialized = False
        self.is_detecting = False
        self.bridge = CvBridge()
        self.pipeline: Optional[DoubleYoloPipeline] = None
        self.latest_image: Optional[np.ndarray] = None
        self.processing_lock = threading.Lock()
        
        # 结果缓存队列
        self.cone_queue = ResultQueue(self.queue_size, "锥桶")
        self.number_queue = ResultQueue(self.queue_size, "数字")
        self.dashboard_queue = ResultQueue(self.queue_size, "仪表盘")
        
        # 检测状态跟踪
        self.detection_start_time = 0
        self.frames_processed_in_task = 0
        self.last_process_time = 0
        
        # 统计信息
        self.total_tasks = 0
        self.last_complete_result = "None-None-None"
        
        # ROS话题
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.status_pub = rospy.Publisher('/yolo_status', String, queue_size=10)
        
        self.command_sub = rospy.Subscriber('/yolo_detect', String, self.command_callback, queue_size=1)
        self.image_sub = rospy.Subscriber('/camera/image_raw', Image, self.image_callback, queue_size=1)
        
        rospy.loginfo("=== 完整目标检测YOLO节点初始化完成 ===")
        rospy.loginfo(f"模型路径: {self.model_path}")
        rospy.loginfo(f"设备: {self.device if self.device else '自动选择'}")
        rospy.loginfo(f"调试模式: {self.debug}")
        rospy.loginfo(f"检测参数: first_conf={self.first_stage_conf}, second_conf={self.second_stage_conf}, final_conf={self.final_conf}")
        rospy.loginfo(f"最大检测时间: {self.max_detection_time}s")
        rospy.loginfo(f"结果队列大小: {self.queue_size}")
        rospy.loginfo("等待yolostart命令开始完整检测任务...")
        
        # 发布初始状态
        self.publish_status("节点已初始化，等待yolostart命令开始完整检测")

    def publish_status(self, message: str):
        """发布状态信息"""
        try:
            status_msg = String()
            status_msg.data = f"[CompleteDetection] {message}"
            self.status_pub.publish(status_msg)
        except Exception as e:
            rospy.logwarn(f"发布状态失败: {e}")

    def initialize_pipeline(self):
        """初始化YOLO管道（懒加载）"""
        if self.pipeline_initialized:
            return True
            
        try:
            rospy.loginfo("正在初始化双阶段YOLO管道...")
            self.publish_status("正在初始化YOLO管道...")
            
            if not os.path.exists(self.model_path):
                error_msg = f"模型文件不存在: {self.model_path}"
                rospy.logerr(error_msg)
                self.publish_status(f"错误: {error_msg}")
                return False
            
            # 初始化管道
            self.pipeline = DoubleYoloPipeline(
                model_path=self.model_path,
                device=self.device,
                debug=self.debug,
                save_dir=self.save_dir
            )
            
            self.pipeline_initialized = True
            rospy.loginfo("✅ YOLO管道初始化完成")
            return True
            
        except Exception as e:
            error_msg = f"管道初始化失败: {e}"
            rospy.logerr(error_msg)
            self.publish_status(f"初始化失败: {e}")
            return False

    def command_callback(self, msg: String):
        """处理控制命令"""
        command = msg.data.strip().lower()
        rospy.loginfo(f"收到命令: {command}")
        
        if command == "yolostart":
            self.start_complete_detection()
        elif command == "yolostop":
            self.stop_detection()
        elif command == "yolostatus":
            self.report_status()
        elif command == "yoloreset":
            self.reset_system()
        else:
            rospy.logwarn(f"未知命令: {command}")
            self.publish_status(f"未知命令: {command}")

    def start_complete_detection(self):
        """开始完整检测任务"""
        if self.is_detecting:
            rospy.logwarn("检测任务正在进行中")
            self.publish_status("检测任务正在进行中，请等待完成")
            return
            
        if self.latest_image is None:
            rospy.logwarn("没有可用的图像数据")
            self.publish_status("错误: 没有可用的图像数据")
            return
            
        # 初始化管道（如果未初始化）
        if not self.initialize_pipeline():
            return
            
        # 重置检测状态和队列
        self.clear_all_queues()
        self.detection_start_time = time.time()
        self.frames_processed_in_task = 0
        self.last_process_time = 0
        self.is_detecting = True
        
        rospy.loginfo("🚀 开始完整检测任务...")
        rospy.loginfo("=" * 80)
        rospy.loginfo("开始逐帧检测，等待三个队列同时填满...")
        self.publish_status("检测任务开始 - 寻找锥形桶、数字、仪表盘...")

    def clear_all_queues(self):
        """清空所有结果队列"""
        self.cone_queue.clear()
        self.number_queue.clear()
        self.dashboard_queue.clear()
        rospy.loginfo("🧹 所有结果队列已清空")

    def stop_detection(self):
        """停止检测任务"""
        if self.is_detecting:
            self.is_detecting = False
            elapsed_time = time.time() - self.detection_start_time
            rospy.loginfo("=" * 80)
            rospy.loginfo(f"🛑 检测任务已手动停止 (耗时: {elapsed_time:.1f}s, 处理帧数: {self.frames_processed_in_task})")
            self.publish_status(f"检测任务已停止 (耗时: {elapsed_time:.1f}s)")
            self.print_queue_status()
        else:
            rospy.loginfo("当前没有进行中的检测任务")
            self.publish_status("当前没有进行中的检测任务")

    def reset_system(self):
        """重置系统"""
        self.stop_detection()
        self.cleanup_pipeline()
        self.clear_all_queues()
        time.sleep(1)
        
        self.total_tasks = 0
        self.last_complete_result = "None-None-None"
        
        rospy.loginfo("🔄 系统已重置")
        self.publish_status("系统已重置，等待yolostart命令")

    def cleanup_pipeline(self):
        """清理管道资源"""
        if self.pipeline:
            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except:
                pass
            self.pipeline = None
            self.pipeline_initialized = False
            gc.collect()

    def report_status(self):
        """报告状态"""
        if self.is_detecting:
            elapsed_time = time.time() - self.detection_start_time
            status = f"检测中 ({elapsed_time:.1f}s)"
        else:
            status = "等待yolostart命令"
            
        report = (f"状态: {status} | "
                 f"已完成任务: {self.total_tasks} 次 | "
                 f"最后结果: {self.last_complete_result}")
        
        rospy.loginfo(report)
        self.publish_status(report)
        
        if self.is_detecting:
            self.print_queue_status()

    def print_queue_status(self):
        """打印队列状态"""
        rospy.loginfo("📊 当前队列状态:")
        rospy.loginfo(f"   {self.cone_queue.get_status()}")
        rospy.loginfo(f"   {self.number_queue.get_status()}")
        rospy.loginfo(f"   {self.dashboard_queue.get_status()}")

    def image_callback(self, msg: Image):
        """处理图像回调"""
        try:
            # 转换图像
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            self.latest_image = cv_image
            
            # 如果正在检测，处理当前帧
            if self.is_detecting:
                current_time = time.time()
                
                # 检查检测超时
                if current_time - self.detection_start_time > self.max_detection_time:
                    self.handle_detection_timeout()
                    return
                
                # 控制处理频率
                if current_time - self.last_process_time < self.process_interval:
                    return
                    
                # 非阻塞处理
                if self.processing_lock.acquire(blocking=False):
                    try:
                        self.process_detection_frame()
                    finally:
                        self.processing_lock.release()
                        
        except Exception as e:
            rospy.logwarn(f"图像处理回调错误: {e}")

    def handle_detection_timeout(self):
        """处理检测超时"""
        self.is_detecting = False
        elapsed_time = time.time() - self.detection_start_time
        
        # 从队列获取最佳结果
        cone_result = self.cone_queue.get_most_common() or 'None'
        number_result = self.number_queue.get_most_common() or 'None'
        dashboard_result = self.dashboard_queue.get_most_common() or 'None'
        
        partial_result = f"{cone_result}-{number_result}-{dashboard_result}"
        
        result_msg = String()
        result_msg.data = partial_result
        self.result_pub.publish(result_msg)
        
        rospy.loginfo("=" * 80)
        rospy.logwarn(f"⏰ 检测任务超时 ({elapsed_time:.1f}s)")
        self.print_queue_status()
        rospy.loginfo(f"📤 发布部分结果: {partial_result}")
        self.publish_status(f"检测超时，部分结果: {partial_result}")

    def process_detection_frame(self):
        """处理检测帧"""
        if not self.is_detecting or self.latest_image is None:
            return
            
        try:
            # 复制当前图像
            current_image = self.latest_image.copy()
            self.last_process_time = time.time()
            self.frames_processed_in_task += 1
            
            # 使用双阶段管道处理图像
            result = self.pipeline.process_image(
                image=current_image,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                second_stage_conf=self.second_stage_conf,
                final_conf=self.final_conf,
                frame_idx=self.frames_processed_in_task
            )
            
            if 'error' in result:
                rospy.logwarn(f"帧处理错误: {result['error']}")
                return
                
            # 输出每帧结果
            self.print_frame_result(result)
            
            # 更新结果队列
            self.update_result_queues(result)
            
            # 检查是否完成
            if self.are_all_queues_full():
                self.complete_detection_task()
                
        except Exception as e:
            rospy.logerr(f"帧处理失败: {e}")

    def print_frame_result(self, result):
        """打印每帧的检测结果"""
        cone_color = result.get('cone_color', 'None')
        number = result.get('number', 'None')
        dashboard_state = result.get('dashboard_state', 'None')
        
        # 格式化输出
        frame_result = f"[Frame {self.frames_processed_in_task:4d}] "
        frame_result += f"锥桶: {cone_color:>8s} | "
        frame_result += f"数字: {number:>8s} | "
        frame_result += f"仪表盘: {dashboard_state:>8s}"
        
        print(frame_result)  # 使用print确保在终端显示

    def update_result_queues(self, result):
        """更新结果队列"""
        cone_color = result.get('cone_color')
        number = result.get('number')
        dashboard_state = result.get('dashboard_state')
        
        # 添加结果到相应队列
        cone_added = self.cone_queue.add(cone_color)
        number_added = self.number_queue.add(number)
        dashboard_added = self.dashboard_queue.add(dashboard_state)
        
        # 如果有新结果添加到队列，显示队列状态
        if cone_added or number_added or dashboard_added:
            queue_status = f"队列状态: 锥桶({len(self.cone_queue.queue)}/{self.queue_size}) "
            queue_status += f"数字({len(self.number_queue.queue)}/{self.queue_size}) "
            queue_status += f"仪表盘({len(self.dashboard_queue.queue)}/{self.queue_size})"
            
            if cone_added:
                queue_status += f" [锥桶+{cone_color}]"
            if number_added:
                queue_status += f" [数字+{number}]"
            if dashboard_added:
                queue_status += f" [仪表盘+{dashboard_state}]"
                
            print(f"    {queue_status}")

    def are_all_queues_full(self):
        """检查所有队列是否都已满"""
        return (self.cone_queue.is_full() and 
                self.number_queue.is_full() and 
                self.dashboard_queue.is_full())

    def complete_detection_task(self):
        """完成检测任务"""
        self.is_detecting = False
        elapsed_time = time.time() - self.detection_start_time
        self.total_tasks += 1
        
        # 从队列获取最终结果
        final_cone = self.cone_queue.get_most_common()
        final_number = self.number_queue.get_most_common()
        final_dashboard = self.dashboard_queue.get_most_common()
        
        # 构建完整结果
        complete_result = f"{final_cone}-{final_number}-{final_dashboard}"
        self.last_complete_result = complete_result
        
        # 发布结果
        result_msg = String()
        result_msg.data = complete_result
        self.result_pub.publish(result_msg)
        
        # 输出完成信息
        print("=" * 80)
        rospy.loginfo(f"🎉 检测任务完成! 最终结果: {complete_result}")
        rospy.loginfo(f"📊 最终队列状态:")
        rospy.loginfo(f"   {self.cone_queue.get_status()}")
        rospy.loginfo(f"   {self.number_queue.get_status()}")
        rospy.loginfo(f"   {self.dashboard_queue.get_status()}")
        rospy.loginfo(f"⏱️  任务统计: 耗时 {elapsed_time:.1f}s, 处理 {self.frames_processed_in_task} 帧, "
                     f"平均 {self.frames_processed_in_task/elapsed_time:.1f} FPS")
        print("=" * 80)
        
        # 清空队列
        self.clear_all_queues()
        
        self.publish_status(f"任务完成: {complete_result} (耗时: {elapsed_time:.1f}s, {self.frames_processed_in_task}帧) - 等待下次yolostart命令")

    def run(self):
        """运行节点"""
        rospy.loginfo("完整目标检测YOLO节点开始运行...")
        self.publish_status("节点运行中，发送yolostart命令开始完整检测任务")
        
        try:
            rospy.spin()
        except KeyboardInterrupt:
            rospy.loginfo("收到停止信号")
        finally:
            self.stop_detection()
            self.cleanup_pipeline()
            rospy.loginfo("完整目标检测YOLO节点已关闭")


if __name__ == '__main__':
    try:
        node = YoloCompleteDetectionNode()
        node.run()
    except rospy.ROSInterruptException:
        pass
    except Exception as e:
        rospy.logerr(f"节点启动失败: {e}")
        sys.exit(1)